import torch
from torch.utils.data import Dataset, DataLoader


def _generate_sin(num=100):
    x = torch.linspace(0, 2 * torch.pi, num)
    y = torch.sin(x) + 2
    y += torch.normal(0, 0.1, y.shape)
    return y


class SinDataset(Dataset):
    def __init__(self, y, seq_len, skip=1):  # seq_len 既为输入长度 也为输出长度
        super().__init__()
        steps = 100 - seq_len + 1 - skip  # 因为取到的最后一个序列没有输出序列
        self.inputs = torch.zeros((steps, seq_len))
        self.outputs = torch.zeros((steps, seq_len))
        # 通过顺序关系，将数据取出
        for i in range(steps):
            self.inputs[i] = y[i:i + seq_len]
            self.outputs[i] = y[i + skip:i + skip + seq_len]

    def __getitem__(self, index):
        return self.inputs[index], self.outputs[index]

    def __len__(self):
        return len(self.inputs)


def generate_sin_loader(num=100, batch_size=10, seq_len=5, skip=1):
    y = _generate_sin(num)
    dataset = SinDataset(y, seq_len, skip)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True,drop_last=True)
    return dataloader
